import torch
import torchvision
import matplotlib.patches as patches
import matplotlib.lines as lines
import matplotlib.pyplot as plt
import cv2 as cv

dev = torch.device('cuda')

model = torchvision.models.detection.keypointrcnn_resnet50_fpn(weights=torchvision.models.detection.KeypointRCNN_ResNet50_FPN_Weights.COCO_V1)
model = model.to(dev)
model.eval()

cap = cv.VideoCapture(2)
cap.set(cv.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv.CAP_PROP_FRAME_HEIGHT, 1080)

plt.ion()

fig, ax = plt.subplots()

dict.map = lambda self, fn : map(fn, self.items())

while True:
    frame = cap.read()[1]
    frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
    frame1 = torch.tensor(frame, device=dev).permute(2, 0, 1) / 255
    prediction = model([frame1])[0]
    map(lambda pr : pr.cpu().detach().numpy(), prediction.values())
    for key, value in prediction.items():
        prediction[key] = prediction[key].cpu().detach().numpy()
    prediction = [
        (box, label, score, keypoint)
        for box, label, score, keypoint
        in zip(
            prediction['boxes'],
            prediction['labels'],
            prediction['scores'],
            prediction['keypoints']
        )
        if score > 0.9
    ]
    ax.clear()
    ax.imshow(frame)
    for box, label, score, keypoint in prediction:
        ax.add_patch(patches.Rectangle(box[0 : 2], *(box[2 : 4] - box[0 : 2]), fill=False, edgecolor='#FFFFFF'))
        ax.text(*box[0 : 2], label, color='#FFFFFF')
        for l in [[3, 1, 0, 5, 7, 9], [4, 2, 0, 6, 8, 10], [0, 11, 13, 15], [0, 12, 14, 16]]:
            ax.add_line(lines.Line2D(*list(map(list, zip(*[[keypoint[k, 0], keypoint[k, 1]] for k in l])))))
    fig.canvas.draw()
    fig.canvas.flush_events()